from enum import Enum
from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel
from app.schemas.metrics_schema import Chart, BarChart, ScatterPlot, LineChart
from app.schemas.facenet_base_schema import TaskStatus, ProductType

# THGRL算法的输入参数
class THGRLInputParams(BaseModel):
    graph_sequence: List[Dict[str, Any]]  # 图序列，格式为[{nodes: [...], edges: [...]}, ...]
    embedding_dim: int = 64  # 嵌入维度
    time_steps: int = 5  # 使用的时间步数
    learning_rate: float = 0.01  # 学习率
    epochs: int = 100  # 训练轮数

# 节点嵌入结果
class NodeEmbedding(BaseModel):
    node_id: str  # 节点ID
    embedding: List[float]  # 节点嵌入向量
    node_type: Optional[str] = None  # 节点类型

# 链接预测结果
class LinkPrediction(BaseModel):
    source_node: str  # 源节点
    target_node: str  # 目标节点
    probability: float  # 链接概率
    edge_type: Optional[str] = None  # 边类型

# THGRL算法的输出参数
class THGRLOutputParams(BaseModel):
    node_embeddings: List[NodeEmbedding]  # 所有节点的嵌入结果
    link_predictions: Optional[List[LinkPrediction]] = None  # 链接预测结果
    algorithm: str = "THGRL"  # 算法名称
    parameters: Dict[str, Any]  # 算法参数
    training_loss: Optional[List[float]] = None  # 训练损失记录

# 扩展输入参数
class InputParams(BaseModel):
    thgrl_params: Optional[THGRLInputParams] = None
    # 其他可能的参数

# 扩展输出参数
class OutputParams(BaseModel):
    thgrl_results: Optional[THGRLOutputParams] = None
    # 其他可能的输出

# 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams

# 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    
    input_params: InputParams
    error_message: Optional[str] = None
    
    metrics: List[Chart] = []

# 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    
    error_message: Optional[str] = None
    
    output_params: OutputParams
    
    metrics: List[Chart] = [] 